package com.dequan.springaiproject.config;

import com.dequan.springaiproject.constant.PromptConstant;
import com.dequan.springaiproject.manager.ChatClientManager;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
import org.springframework.ai.chat.memory.repository.jdbc.JdbcChatMemoryRepository;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.ai.openai.OpenAiChatOptions;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.util.Assert;

import java.util.Collections;
import java.util.List;

/**
 * Qwen 模型配置类
 * 1. 如何在 chatModel 的基础上上创建 chatClient
 * 2. 如何创建一个模型下不同身份的 ChatClient 用于不同场景的问答
 * 3. 做水平扩展的时候，相同的大模型服务请求要做到负载均衡，可以使用Nginx
 * 如果是使用不同的代理商的提供的相同的大模型服务，如何轮训不同的服务结点？
 * 当某个代理商挂了呢？怎么动态修改服务结点List？
 *
 *
 */
@Configuration
public class QwenModelConfiguration {

    /**
     * spring AI 在引入依赖后会自动配置 ChatModel 和 ChatModel.Builder
     * 但不会自动创建 ChatClient Bean
     * @param chatModel
     * @return
     */
    @Bean
    public ChatClient chatClient(@Autowired ChatModel chatModel) {
        return ChatClient.create(chatModel);
    }

    /**
     * 创建一个医生角色的 ChatClient
     * @param chatClientBuilder
     * @return
     */
    @Bean
    public ChatClient doctorChatClient(ChatClient.Builder chatClientBuilder) {
        return chatClientBuilder
                .defaultSystem(PromptConstant.DOCTOR_PROMPT)
                .defaultOptions(OpenAiChatOptions.builder()
                        .model("qwen-max")
                        .build())
                .build();
    }

    /**
     * 创建一个学生角色的 ChatClient
     * @param chatClientBuilder
     * @return
     */
    @Bean
    public ChatClient studentChatClient(ChatClient.Builder chatClientBuilder) {
        return chatClientBuilder
                .defaultSystem(PromptConstant.STUDENT_PROMPT)
                .defaultOptions(OpenAiChatOptions.builder()
                        .temperature(3.0)
                        .build())
                .build();
    }


    /**
     * 创建并配置ChatClientManager Bean实例
     *
     * @param chatModel 用于构建聊天客户端的基础聊天模型
     * @param openAiApi OpenAI API配置实例，用于构建不同的API连接
     * @param baseUrls OpenAI服务的基础URL列表，支持多个URL配置
     * @param apiKeys 对应的API密钥列表，与baseUrls一一对应
     * @return 配置完成的ChatClientManager实例，管理多个聊天客户端
     */
    @Bean
    public ChatClientManager chatClientManager(
            @Autowired OpenAiChatModel chatModel,
            @Autowired OpenAiApi openAiApi,
            @Value("${spring.ai.openai.base-url}") List<String> baseUrls,
            @Value("${spring.ai.openai.api-key}") List<String> apiKeys
            ) {
        int n = baseUrls.size();
        int m = apiKeys.size();
        Assert.isTrue(n == m, "baseUrls and apiKeys must have the same size");

        if (n == 0) {
            return new ChatClientManager(Collections.emptyList());
        }

        ChatClientManager.Builder builder = new ChatClientManager.Builder();
        for (int i = 0; i < baseUrls.size(); i++) {
            ChatModel currentChatModel = chatModel.mutate()
                    .defaultOptions(OpenAiChatOptions.builder()
                            .model("qwen-plus")
                            .temperature(0.2 + i * 0.5) // TODO 临时解决不同模型参数的配置
                            .build())
                    .openAiApi(openAiApi.mutate()
                            .baseUrl(baseUrls.get(i))
                            .apiKey(apiKeys.get(i))
                            .build())
                    .build();
            builder.addChatClient(ChatClient.builder(currentChatModel).build());
        }
        return builder.build();
    }

    /*
     * 创建一个带数据库实现的 ChatMemory 的 ChatClient
     */
    @Bean
    public ChatClient withMemoryAdvisorChatClient(
            ChatClient.Builder chatClientBuilder,
            ChatMemory chatMemory
    ) {
        // 创建一个 ChatClient
        return chatClientBuilder
                .defaultAdvisors(
                        MessageChatMemoryAdvisor.builder(chatMemory).build()
                ).build();
    }

    /*
     * 创建一个带数据库实现的 ChatMemory 的 ChatClient
     */
    @Bean
    public ChatClient withDBMemoryAdvisorChatClient (
            ChatClient.Builder chatClientBuilder,
            JdbcChatMemoryRepository jdbcChatMemoryRepository
    ) {
        MessageWindowChatMemory build = MessageWindowChatMemory.builder()
                .chatMemoryRepository(jdbcChatMemoryRepository)
                .maxMessages(10)
                .build();

        return chatClientBuilder
                .defaultAdvisors(
                        MessageChatMemoryAdvisor.builder(build).build()
                ).build();
    }
}